package com.cat.binaryTree;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;

/**
 * @author 曲大人的喵
 * @description https://leetcode.cn/problems/number-of-good-leaf-nodes-pairs/
 * @create 2025/9/13 16:49
 * @since JDK17
 */

public class Solution29 {
    int ans, distance;
    List<Integer> f(TreeNode root) {
        List<Integer> distances = new ArrayList<>();
        if (root == null) {
            return distances;
        }
        if (root.left == null && root.right == null) {
            distances.add(0);   // 距离值
            return distances;
        }
        var left = f(root.left);
        var right = f(root.right);
        for (int i = 0; i < left.size(); i++) {
            int ld = left.get(i) + 1;
            left.set(i, ld);
            if (ld <= distance) {
                distances.add(ld);
            }
        }
        for (int i = 0; i < right.size(); i++) {
            int rd = right.get(i) + 1;
            right.set(i, rd);
            if (rd <= distance) {
                distances.add(rd);
            }
        }
        for (int i = 0; i < left.size(); i++) {
            for (int j = 0; j < right.size(); j++) {
                ans += left.get(i) + right.get(j) <= distance ? 1 : 0;
            }
        }
        return distances;
    }

    public int countPairs(TreeNode root, int distance) {
        this.distance = distance;
        f(root);
        return ans;
    }
}
